import torch
import json
import os
from datasets import load_dataset, Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model

# 启用MPS GPU加速（M4 Pro优化）
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

# 模型配置 - 使用更小的模型适配M4 Pro内存
MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)

def load_local_dataset(data_path):
    """
    加载本地数据集
    支持的格式：
    1. JSON文件：包含instruction, input(可选), output字段的列表
    2. JSONL文件：每行一个JSON对象
    
    Args:
        data_path (str): 数据文件路径
    
    Returns:
        Dataset: HuggingFace Dataset对象
    """
    if not os.path.exists(data_path):
        raise FileNotFoundError(f"数据文件不存在: {data_path}")
    
    data = []
    
    if data_path.endswith('.json'):
        # 加载JSON文件
        with open(data_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    elif data_path.endswith('.jsonl'):
        # 加载JSONL文件
        with open(data_path, 'r', encoding='utf-8') as f:
            for line in f:
                data.append(json.loads(line.strip()))
    else:
        raise ValueError("不支持的文件格式，请使用.json或.jsonl文件")
    
    # 验证数据格式
    required_fields = ['instruction', 'output']
    for i, item in enumerate(data):
        for field in required_fields:
            if field not in item:
                raise ValueError(f"第{i+1}行数据缺少必需字段: {field}")
    
    return Dataset.from_list(data)

# 数据集配置
# 只使用本地数据集
LOCAL_DATA_PATH = './data/train_data.json'  # 本地数据文件路径

# 加载本地数据集
print(f"正在加载本地数据集: {LOCAL_DATA_PATH}")
train_dataset = load_local_dataset(LOCAL_DATA_PATH)
dataset = {'train': train_dataset}

# 格式化函数：把指令/输入/输出拼成 prompt
def format_example(example):
    if example.get("input"):
        full_text = f"指令: {example['instruction']}\n输入: {example['input']}\n回答: {example['output']}"
    else:
        full_text = f"指令: {example['instruction']}\n回答: {example['output']}"
    
    # 对完整文本进行分词
    tokenized = tokenizer(
        full_text,
        truncation=True,
        padding=False,
        max_length=512,
    )
    
    return {
        "input_ids": tokenized["input_ids"],
        "labels": tokenized["input_ids"].copy(),  # 标签与输入相同（自回归训练）
    }

# 数据预处理
# 本地数据集处理
tokenized_dataset = dataset["train"].map(format_example, remove_columns=dataset["train"].column_names)

# 加载模型（Mac M4 Pro优化配置）
print("正在加载模型...")
    # MPS GPU配置 - 避免meta tensor问题
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,  # 使用float32确保MPS兼容性
    device_map=None,            # 禁用自动设备映射避免meta tensor
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
    # 手动移动模型到MPS设备
model = model.to(device)

print(f"模型已加载到设备: {device}")

# LoRA配置（极简配置）
lora_config = LoraConfig(
    r=16,                    # 增加LoRA rank提高表达能力
    lora_alpha=32,          # 相应增加alpha
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 训练更多attention层
    lora_dropout=0.0,       # 适当的dropout
    bias="none",            # 不训练bias
    task_type="CAUSAL_LM"   # 因果语言模型
)

model = get_peft_model(model, lora_config)

# 确保LoRA模型也在正确的设备上
model = model.to(device)
print(f"LoRA模型已移动到设备: {device}")

# 打印可训练参数信息
model.print_trainable_parameters()

# 检查是否有可训练参数
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
all_params = sum(p.numel() for p in model.parameters())
print(f"\n=== 参数统计 ===")
print(f"可训练参数数量: {trainable_params:,}")
print(f"总参数数量: {all_params:,}")
print(f"可训练参数比例: {100 * trainable_params / all_params:.2f}%")

# 确保模型在正确的设备上并且处于训练模式
model = model.to(device)
model.train()
print(f"模型已设置为训练模式并移动到设备: {device}")

# 训练参数（Mac M4 Pro GPU优化配置）
training_args = TrainingArguments(
    per_device_train_batch_size=1,       # GPU可以支持更大batch
    gradient_accumulation_steps=1,       # 适中的梯度累积
    # warmup_steps=10,                     # 适当的warmup
    # max_steps=50,    
    num_train_epochs=10,
    warmup_ratio=0.1,
    learning_rate=3e-4,                  # 降低学习率避免不稳定
    fp16=False,                          # MPS暂不支持fp16混合精度
    bf16=False,                          # MPS暂不支持bf16
    logging_steps=2,
    output_dir="./qwen15b-lora",
    # save_strategy="steps",
    # save_steps=100,
    save_strategy="epoch",
    dataloader_pin_memory=False,         # Mac优化
    remove_unused_columns=False,
    optim="adamw_torch",
    gradient_checkpointing=False,        # 禁用梯度检查点避免MPS兼容性问题
    dataloader_num_workers=0,
    report_to=[],
    eval_strategy="no",
    no_cuda=True
)
   

# 准备训练数据
# 本地数据集处理
train_data = tokenized_dataset.shuffle().select(range(min(1000, len(tokenized_dataset))))
# 数据整理器，处理动态padding
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # 不使用掩码语言模型
    pad_to_multiple_of=8,  # 填充到8的倍数
)

# 创建训练器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    data_collator=data_collator,
)

# 训练前验证
print("\n=== 训练前验证 ===")
print(f"训练数据集大小: {len(train_data)}")
print(f"设备类型: {device}")
print(f"模型设备: {next(model.parameters()).device}")
print(f"模型训练模式: {model.training}")

# 验证第一个batch
try:
    sample_batch = next(iter(trainer.get_train_dataloader()))
    print(f"数据加载器正常，batch大小: {len(sample_batch['input_ids'])}")
    print(f"输入序列长度: {sample_batch['input_ids'].shape}")
except Exception as e:
    print(f"数据加载器错误: {e}")

# 开始训练
print("\n=== 开始训练 ===")
try:
    trainer.train()
    print("训练完成")
except Exception as e:
    print(f"训练过程中出现错误: {e}")
    raise

# 保存微调好的 LoRA 权重
model.save_pretrained("./models/qwen15b-lora")
tokenizer.save_pretrained("./models/qwen15b-lora")
